3.3 하드웨어 인식(Hardware-Aware) 알고리즘과 GPU 메모리 계층 최적화

현대 인공지능 연구의 최전선에서 모델 아키텍처의 혁신은 더 이상 수학적 이론의 영역에만 머물지 않는다. 트랜스포머(Transformer)가 지난 수년간 자연어 처리(NLP)를 비롯한 딥러닝 분야를 지배할 수 있었던 핵심적인 이유는 어텐션(Attention) 메커니즘이 GPU와 같은 대규모 병렬 처리 하드웨어에 적합했기 때문이다. 그러나 시퀀스 길이(L)에 대해 2차원적(O(L^2))으로 증가하는 연산 및 메모리 복잡도는 긴 문맥(Long Context)을 처리하는 데 있어 넘을 수 없는 물리적 장벽으로 작용해 왔다. Mamba 아키텍처의 등장이 학계와 산업계에 충격을 준 이유는 단순히 선형 시간 복잡도(O(L))를 달성한 이론적 모델(State Space Model, SSM)을 제안했기 때문이 아니라, 이를 실제 하드웨어의 물리적 특성에 맞춰 극한으로 최적화한 하드웨어 인식(Hardware-Aware) 알고리즘을 구현했기 때문이다.1

본 장에서는 Mamba가 어떻게 이론적인 선택적 상태 공간 모델(Selective SSM)을 GPU의 메모리 계층 구조(Memory Hierarchy)와 결합하여 트랜스포머를 능가하는 추론 속도와 학습 효율성을 달성했는지 심층적으로 분석한다. 특히 컴퓨트 바운드(Compute-bound)와 메모리 바운드(Memory-bound)의 개념부터 시작하여, 커널 융합(Kernel Fusion), 병렬 스캔(Parallel Scan), 그리고 재계산(Recomputation) 전략에 이르는 기술적 세부 사항을 엔지니어링 관점에서 상세히 기술한다.

1. 딥러닝 연산의 물리적 병목: 메모리 장벽(Memory Wall)의 이해

Mamba의 설계 철학을 이해하기 위해서는 먼저 현대 가속기(Accelerator)가 직면한 근본적인 병목 현상을 이해해야 한다. 딥러닝 연산은 크게 두 가지 유형으로 분류된다. 첫째는 행렬 곱셈(Matrix Multiplication)이나 고밀도 합성곱(Convolution)과 같이 프로세서의 산술 논리 장치(ALU) 처리 능력이 속도를 결정하는 컴퓨트 바운드(Compute-bound) 작업이다. 둘째는 정규화(Normalization), 요소별(Element-wise) 연산, 그리고 SSM의 스캔(Scan) 연산과 같이 데이터를 메모리에서 레지스터로 불러오거나 저장하는 속도가 전체 성능을 제한하는 메모리 바운드(Memory-bound) 작업이다.1

GPU 제조사들은 텐서 코어(Tensor Core)와 같은 특수 연산 유닛을 통해 연산 능력(FLOPS)을 비약적으로 향상시켜 왔으나, 메모리 대역폭(Bandwidth)의 발전 속도는 이에 미치지 못했다. 이러한 연산 속도와 메모리 속도 간의 격차 확대, 즉 메모리 장벽(Memory Wall) 문제는 긴 시퀀스를 처리하는 모델에서 더욱 두드러진다.

1.1 GPU 메모리 계층 구조와 대역폭 격차 분석

현대 데이터센터 GPU인 NVIDIA A100과 H100의 메모리 계층 구조는 용량은 크지만 상대적으로 느린 **HBM(High Bandwidth Memory)**과, 용량은 극히 작지만 매우 빠른 온칩(On-chip) 메모리인 **SRAM(Static Random Access Memory)**으로 구성된다. Mamba의 알고리즘은 이 두 메모리 계층 간의 속도 차이를 활용하는 데 초점을 맞춘다.1

아래 표는 주요 GPU 아키텍처의 메모리 사양을 비교한 것이다. HBM과 SRAM(L1/L2 캐시 포함) 간의 대역폭 차이가 수십 배에 달함을 확인할 수 있다.

트랜스포머의 어텐션 메커니즘은 Q, K, V 행렬을 HBM에서 반복적으로 읽어오는 과정에서 메모리 대역폭의 한계에 부딪힌다. 특히 시퀀스 길이가 길어질수록 어텐션 맵(L \times L)의 크기가 폭발적으로 증가하여 HBM 대역폭을 포화시킨다. 반면, 기존의 순환신경망(RNN)이나 SSM은 순차적인 데이터 의존성으로 인해 병렬 처리가 어렵고, 각 타임스텝마다 중간 상태(Hidden State)를 HBM에 쓰고 다시 읽어오는 잦은 I/O 요청(Memory Access Overhead)을 발생시켜 GPU의 연산 유닛을 유휴 상태(Idle)로 만든다.2

Mamba의 핵심 통찰은 **“메모리 바운드 작업에서 HBM과 SRAM 사이의 데이터 이동(I/O)을 최소화해야 한다”**는 것이다.1 이를 위해 Mamba는 커널 융합(Kernel Fusion) 기술을 사용하여 주요 연산을 SRAM 내부에서 처리하고, HBM 접근을 입력과 최종 출력 단계로 제한하는 전략을 취한다.

2. 선택적 스캔(Selective Scan)과 하드웨어 인식 상태 확장

Mamba의 가장 큰 특징인 ‘선택적(Selective)’ 메커니즘은 입력 데이터의 내용에 따라 모델의 파라미터 (\Delta, B, C)가 동적으로 변하는 것이다. 이는 기존의 LTI(Linear Time Invariant) SSM이 누렸던 효율적인 합성곱(Convolution) 연산(O(L \log L))과 고속 푸리에 변환(FFT)의 이점을 포기하게 만든다.14 대신 순환(Recurrent) 모드를 사용해야 하는데, 단순한 순환 연산은 O(L)의 선형 복잡도를 가지지만 병렬화가 어렵다는 치명적인 단점이 있다.

이 딜레마를 해결하기 위해 Mamba는 **하드웨어 인식 상태 확장(Hardware-Aware State Expansion)**이라는 개념을 도입하여 알고리즘을 재설계했다.1

2.1 상태 구체화(State Materialization)의 제어 전략

SSM의 잠재 상태(Latent State) h는 차원이 (B, L, D, N)으로 매우 크다. 여기서 B는 배치 크기, L은 시퀀스 길이, D는 채널 수, N은 상태 차원이다. 일반적인 딥러닝 프레임워크(PyTorch 등)로 이를 구현하면, 모든 중간 상태 h를 HBM에 구체화(Materialize)하여 저장하게 된다. 이는 막대한 메모리 사용량과 I/O 대역폭 소모를 유발한다.1

Mamba는 이 확장된 상태 h를 HBM에 절대 구체화하지 않는다.1 대신, 다음과 같은 최적화된 데이터 흐름을 따른다:

1. 파라미터 로드 (HBM \to SRAM): SSM의 파라미터 (\Delta, A, B, C)와 입력 x를 느린 HBM에서 빠른 SRAM으로 로드한다. 이때 파라미터의 크기는 (B, L, D) 또는 (B, L, N) 수준으로, 확장된 상태 h (B, L, D, N)보다 훨씬 작다.

2. SRAM 내 연산 집중: SRAM 내부에서 이산화(Discretization)와 순환(Recurrence) 연산을 수행한다.

• 이산화: 연속 시간 파라미터를 이산 시간 파라미터로 변환한다.
  \bar{A} = \exp(\Delta A)

  \bar{B} = (\Delta A)^{-1}(\exp(\Delta A) - I) \cdot \Delta B

• 이 과정에서 생성되는 중간 상태 h는 SRAM 내의 레지스터나 공유 메모리에만 잠시 존재하며, 다음 연산에 즉시 사용된 후 폐기되거나 덮어써진다.3

3. 결과 저장 (SRAM \to HBM): 최종 출력 y (크기 (B, L, D))만을 계산하여 HBM에 다시 기록한다.

이러한 방식은 메모리 대역폭 요구량을 O(BLDN)에서 O(BL(D+N))으로 획기적으로 줄여준다. 이는 FlashAttention이 어텐션 행렬을 HBM에 기록하지 않고 SRAM 내에서 타일링(Tiling)을 통해 처리하는 원리와 유사하며, Mamba를 **“순환 모델을 위한 FlashAttention”**이라고 부르는 이유이기도 하다.17

3. 커널 융합(Kernel Fusion) 아키텍처의 구현

Mamba의 하드웨어 가속은 **커널 융합(Kernel Fusion)**을 통해 완성된다. 커널 융합은 여러 개의 작은 GPU 커널(함수)을 하나의 큰 커널로 합쳐, 메모리 I/O를 줄이고 커널 실행 오버헤드를 감소시키는 기법이다.13

Mamba의 선택적 스캔 연산은 구체적으로 다음과 같은 단계들이 하나의 커널로 융합되어 실행된다:

1. 입력 투영 및 로딩: 입력 토큰 x_t와 파라미터 \Delta, A, B, C를 SRAM의 제한된 용량에 맞춰 블록 단위로 로드한다.
2. 이산화 및 게이팅: 로드된 파라미터를 이용해 \bar{A}, \bar{B}를 계산한다. 이 단계는 메모리 집약적인 작업이므로 SRAM 내에서 수행하여 지연 시간을 최소화한다.
3. 병렬 스캔 (Associative Scan): 이산화된 파라미터를 바탕으로 병렬 스캔 알고리즘을 수행하여 상태 h_t를 계산한다. (상세 내용은 3.3.4절 참조)
4. 출력 투영: 계산된 상태 h_t와 C를 곱하여 최종 출력 y_t를 생성한다.

일반적인 구현에서는 각 단계마다 HBM에 데이터를 쓰고 읽어야 했다.

• 융합 전 (Standard Implementation):

Load \to SRAM \to Discretize \to HBM 저장 \to Load \to SRAM \to Scan \to HBM 저장 \to Load \to SRAM \to Multiply…

• 융합 후 (Fused Mamba Kernel):

Load \to SRAM \to (Discretize + Scan + Multiply) \to SRAM \to HBM 저장

이러한 커널 융합은 GPU의 메모리 계층 구조를 최대한 활용하여, 데이터가 가장 빠른 메모리 영역(SRAM/Register)에 머무르는 시간을 극대화한다. 결과적으로 Mamba는 메모리 대역폭이 아닌 연산 능력에 의해 성능이 결정되는 컴퓨트 바운드 영역으로 작업의 성격을 전환시키는 데 성공했다.13

4. 병렬 스캔 알고리즘(Parallel Scan Algorithm)의 수학적 원리

커널 융합이 메모리 대역폭 문제를 해결했다면, 병렬 스캔 알고리즘은 순환 신경망(RNN)의 본질적인 한계인 ’순차적 의존성’을 극복하여 대규모 병렬 처리를 가능하게 했다.1 전통적인 RNN은 시간 t의 상태를 계산하기 위해 t-1의 결과가 필수적이므로, GPU의 수천 개 코어를 동시에 활용할 수 없었다.

Mamba는 **Blelloch(1990)**이 제안한 병렬 접두사 합(Parallel Prefix Sum) 또는 스캔 알고리즘을 선택적 SSM에 적용했다.21 이 알고리즘은 순차적으로 O(L) 시간이 걸리는 작업을 O(\log L) 시간 복잡도로 단축시킨다.

4.1 결합 법칙(Associativity)의 활용

SSM의 상태 업데이트 식 h_t = \bar{A}_t h_{t-1} + \bar{B}_t x_t는 1차 선형 점화식 형태이다. 이를 연산자 형태로 추상화하면 결합 법칙(Associative Property)이 성립함을 알 수 있다.20
(O_j \circ O_i) \circ O_k = O_j \circ (O_i \circ O_k)
결합 법칙이 성립한다는 것은 연산의 순서를 변경해도 결과가 같다는 것을 의미하며, 이는 시퀀스를 여러 청크(Chunk)로 나누어 병렬로 처리한 뒤 합칠 수 있음을 시사한다.

4.2 Blelloch 알고리즘의 단계

Mamba의 스캔 연산은 트리 기반의 환원(Reduction) 구조를 따른다 22:

1. 업스윕(Up-sweep) 또는 리듀스(Reduce) 단계: 인접한 요소들을 쌍으로 묶어 연산 결과를 상위 노드로 올린다. 예를 들어, t=1,2를 묶고, t=3,4를 묶는 식이다. 이 과정은 병렬로 수행되며 트리의 높이인 \log_2 L 단계만에 전체 구간의 요약 정보를 계산할 수 있다.
2. 다운스윕(Down-sweep) 단계: 루트 노드에서부터 아래로 내려오며 각 시점의 누적 값(Prefix Sum)을 계산한다. 이 단계 역시 \log_2 L 단계가 소요된다.

이 알고리즘을 통해 Mamba는 길이 L인 시퀀스를 처리할 때, p개의 병렬 프로세서를 사용하여 O(L/p + \log p)의 시간 복잡도를 달성한다. L이 충분히 크고 p가 많을 경우(현대 GPU 상황), 이는 사실상 O(\log L)에 수렴하여 시퀀스 길이에 대해 극도로 효율적인 확장이 가능하다.21

5. 구현 심층 분석: AWS Neuron 및 NKI 사례 연구

Mamba의 하드웨어 최적화는 NVIDIA GPU(CUDA)에만 국한되지 않는다. AWS의 자체 개발 AI 가속기인 Trainium과 Inferentia(NeuronCore)에서의 구현 사례는 Mamba의 알고리즘이 하드웨어 특성에 맞춰 어떻게 튜닝되어야 하는지를 명확하게 보여준다. AWS의 NKI(Neuron Kernel Interface)를 이용한 Mamba 구현은 타일링(Tiling)과 루프 재정렬(Loop Reordering)의 중요성을 강조한다.25

5.1 타일링(Tiling)과 SBUF 최적화

긴 시퀀스(예: seq_len = 8192 이상)를 처리할 때, 중간 텐서들(delta, u, scanC_accum 등)의 크기는 온칩 스크래치패드 메모리(SBUF)의 용량을 초과할 수 있다. 이를 방치하면 데이터가 느린 오프칩 메모리로 스필링(Spilling)되어 성능이 급격히 저하된다.

이를 방지하기 위해 AWS Neuron 구현에서는 다음과 같은 전략을 사용한다:

1. 시퀀스 차원의 내부 루프 배치: seq_len을 루프의 가장 안쪽(Inner-most loop)에 배치하여 [batch_size, n_channel_tile, state_size, seq_len] 순서로 연산을 수행한다.
2. 타일 기반 처리: 시퀀스를 하드웨어 타일 크기에 맞춰 분할하고, scan_init 변수를 사용해 이전 타일의 스캔 결과를 다음 타일로 전달한다. 이는 루프 간 의존성(Loop-carried dependencies)을 관리하면서도 온칩 메모리 효율을 극대화한다.

5.2 명령어 수준 최적화

• 요소별 곱셈 (Step 3): delta, B, u의 곱셈 연산 시, 메모리 레이아웃을 전치(Transpose)하는 비효율을 피하기 위해 channels 차원을 파티션 차원(Partition Dimension)으로 설정한다. B 텐서는 broadcast_to 명령어를 통해 필요한 차원으로 확장하여 연산 효율을 높인다.
• 결합 스캔 (Step 4): nisa.tensor_tensor_scan과 같은 하드웨어 가속 명령어를 사용하여 스캔 연산을 수행한다.

이러한 최적화 결과, AWS NeuronCore에서 Mamba 구현은 레이턴시를 기존 대비 48% 감소시키고(약 27.8ms), 벡터 엔진 가동률(VectorE activity rate)을 94.85%까지 끌어올리는 성과를 달성했다. 이는 Mamba의 하드웨어 인식 알고리즘이 특정 벤더에 종속되지 않는 보편적인 고성능 컴퓨팅 원리에 기반하고 있음을 시사한다.25

6. 재계산(Recomputation) 전략과 무한 문맥(Infinite Context)

학습(Training) 단계에서는 순전파(Forward pass)뿐만 아니라 역전파(Backward pass)를 위한 메모리 관리도 필수적이다. 역전파를 수행하기 위해서는 순전파 때 계산된 중간 활성화(Activation) 값들이 필요하다. 긴 시퀀스를 처리할 때 이 모든 중간 값을 HBM에 저장해두면 GPU 메모리(VRAM)가 순식간에 고갈된다 (OOM: Out Of Memory). Mamba는 이를 위해 재계산(Recomputation 또는 Gradient Checkpointing) 기법을 하드웨어 레벨에서 공격적으로 활용한다.26

6.0.1 선택적 재계산 (Selective Recomputation)

• 기존 방식: 순전파 시 생성된 크기 (B, L, D, N)의 모든 중간 상태 h를 HBM에 저장한다. 메모리 사용량은 O(BLDN)으로 증가한다.
• Mamba의 방식: 중간 상태 h를 HBM에 저장하지 않고 버린다. 대신 역전파 시, HBM에 저장된 입력 x와 파라미터 (\Delta, A, B, C)를 다시 SRAM으로 불러와서 **필요한 시점의 중간 상태 h를 SRAM 내에서 매우 빠르게 다시 계산(Recompute)**한다.1

이는 “계산 비용이 메모리 I/O 비용보다 저렴하다“는 현대 가속기의 특성을 철저히 이용한 전략이다. HBM에서 거대한 데이터를 읽어오는 시간보다, 작은 입력 데이터를 읽어와서 텐서 코어로 빠르게 다시 계산하는 것이 전체 수행 시간 면에서도 유리하고, 메모리 사용량 측면에서는 비교할 수 없을 만큼 효율적이다. 이 전략 덕분에 Mamba의 메모리 요구량은 시퀀스 길이에 대해 선형적으로 증가하며, FlashAttention이 적용된 트랜스포머와 유사한 수준의 메모리 효율성을 달성한다.28

7. 종합적 성능 평가 및 시사점

3.3절에서 살펴본 Mamba의 하드웨어 인식 알고리즘은 이론적 복잡도 감소를 넘어 실제 시스템 성능에서 괄목할 만한 성과를 이끌어냈다.

1. 추론 속도(Inference Speed): Mamba는 트랜스포머 대비 최대 5배 높은 처리량(Throughput)을 달성했다.1 이는 KV 캐시(Key-Value Cache)가 필요 없는 SSM의 특성과 메모리 계층 최적화가 결합된 결과로, 메모리 대역폭이 제한된 엣지 디바이스나 실시간 응답이 필요한 서비스에서 결정적인 이점을 제공한다.
2. 선형 확장성(Linear Scaling): 시퀀스 길이가 100만(1M) 토큰에 달하는 데이터에 대해서도 성능 저하 없이 선형적으로 확장 가능하다.1 이는 DNA 염기서열 분석, 고해상도 오디오 처리, 장편 소설 분석과 같은 초장문(Long-context) 도메인에서 Mamba가 트랜스포머를 대체할 수 있는 강력한 근거가 된다.

결론적으로 Mamba의 하드웨어 최적화는 알고리즘 설계자가 더 이상 하드웨어를 추상화된 블랙박스로 취급해서는 안 된다는 사실을 증명한다. HBM과 SRAM 사이의 데이터 이동 비용, 커널 융합을 통한 레이턴시 은닉(Latency Hiding), 그리고 병렬 스캔을 통한 순차적 의존성 해결은 포스트 트랜스포머 시대를 여는 핵심 열쇠이다. Mamba는 모델의 표현력(Modeling Power)과 시스템 효율성(System Efficiency)이 공진화(Co-evolution)해야 함을 보여주는 기념비적인 아키텍처라 할 수 있다.

8. 참고 자료

1. Mamba: Linear-Time Sequence Modeling with Selective State Spaces - arXiv, https://arxiv.org/html/2312.00752v2
2. [2312.00752] Mamba: Linear-Time Sequence Modeling with Selective State Spaces - arXiv, https://arxiv.org/abs/2312.00752
3. Mamba Simplified - Part 2 - S4 and Mamba - Prem AI Blog, https://blog.premai.io/s4-and-mamba/
4. Mamba: SSM, Theory, and Implementation in Keras and TensorFlow, https://towardsdatascience.com/mamba-ssm-theory-and-implementation-in-keras-and-tensorflow-32d6d4b32546/
5. Comparing NVIDIA H100 vs A100 GPUs for AI Workloads | OpenMetal IaaS, https://openmetal.io/resources/blog/nvidia-h100-vs-a100-gpu-comparison/
6. NVIDIA A100 | Tensor Core GPU, https://www.nvidia.com/content/dam/en-zz/Solutions/Data-Center/a100/pdf/nvidia-a100-datasheet-nvidia-us-2188504-web.pdf
7. NVIDIA A100 Tensor Core GPU Architecture, https://images.nvidia.com/aem-dam/en-zz/Solutions/data-center/nvidia-ampere-architecture-whitepaper.pdf
8. NVIDIA A100 GPU:, https://hc32.hotchips.org/assets/program/conference/day1/HotChips2020_GPU_NVIDIA_Choquette_v01.pdf
9. INSIDE THE NVIDIA AMPERE A100 GPU IN THETAGPU AND PERLMUTTER JULY 28, 2021, https://www.alcf.anl.gov/sites/default/files/2021-07/ALCF_A100_20210728%5B80%5D.pdf
10. NVIDIA Hopper Architecture In-Depth | NVIDIA Technical Blog, https://developer.nvidia.com/blog/nvidia-hopper-architecture-in-depth/
11. NVIDIA H100 Tensor Core GPU Architecture, https://www.hpctech.co.jp/assets/images/info/catalog/pdf/gtc22-whitepaper-hopper_v1.02.pdf
12. NVIDIA H100 - Newest Data Center GPU built on Hopper Architecture - Exxact Corp., https://www.exxactcorp.com/blog/HPC/NVIDIA-H100
13. A Visual Guide to Mamba and State Space Models - Maarten Grootendorst, https://www.maartengrootendorst.com/blog/mamba/
14. Here Comes Mamba: The Selective State Space Model | Towards Data Science, https://towardsdatascience.com/here-comes-mamba-the-selective-state-space-model-435e5d17a451/
15. Mamba: Linear-Time Sequence Modeling with Selective State Spaces - arXiv, https://arxiv.org/pdf/2312.00752
16. Mamba: Revolutionizing Sequence Modeling with Selective State Spaces - Medium, https://medium.com/@joeajiteshvarun/mamba-revolutionizing-sequence-modeling-with-selective-state-spaces-8a691319b34b
17. state-spaces/mamba: Mamba SSM architecture - GitHub, https://github.com/state-spaces/mamba
18. What Is A Mamba Model? | IBM, https://www.ibm.com/think/topics/mamba-model
19. Mamba architecture : A Leap Forward in Sequence Modeling | by Puneet Hegde - Medium, https://medium.com/@puneetthegde22/mamba-architecture-a-leap-forward-in-sequence-modeling-370dfcbfe44a
20. Mamba No. 5 (A Little Bit Of…) - Sparse Notes, https://jameschen.io/jekyll/update/2024/02/12/mamba.html
21. Recurrent Reinforcement Learning with Memoroids - NIPS papers, https://proceedings.neurips.cc/paper_files/paper/2024/file/19f7f755908372efb25826d61959cdf9-Paper-Conference.pdf
22. Sequential-Parallel Duality in Prefix-Scannable Models - arXiv, https://arxiv.org/html/2506.10918v1
23. MAMBA and State Space Models Explained | by Astarag Mohapatra - Medium, https://athekunal.medium.com/mamba-and-state-space-models-explained-b1bf3cb3bb77
24. Chapter 39. Parallel Prefix Sum (Scan) with CUDA - NVIDIA Developer, https://developer.nvidia.com/gpugems/gpugems3/part-vi-gpu-computing/chapter-39-parallel-prefix-sum-scan-cuda
25. Fused Mamba — AWS Neuron Documentation, https://awsdocs-neuron.readthedocs-hosted.com/en/latest/general/nki/tutorials/fused_mamba.html
26. Mamba: An SSM Method for Efficient and Powerful Sequence Modeling - Medium, https://medium.com/@akdemir_bahadir/mamba-an-ssm-method-for-efficient-and-powerful-sequence-modeling-5dec8f1c849b
27. Activation Recomputation — Megatron Bridge - NVIDIA Documentation, https://docs.nvidia.com/nemo/megatron-bridge/0.1.0/training/activation-recomputation.html
28. Mamba: A Potential Transformer Replacement - Zilliz Learn, https://zilliz.com/learn/mamba-architecture-potential-transformer-replacement